'''
Created on 17 fevr. 2011

@author: paraita
'''
import matplotlib.pyplot as plt 
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import numpy as np
import math
import Image
from PIL import ImageMath

SZ = 64
im = np.zeros((SZ,SZ), np.uint8)
X = range(-SZ/2,SZ/2)
Y = range(-SZ/2,SZ/2)

fp = 4.0
fr = fp/SZ

for i in X:
    for j in Y:
        im[i,j] = 128 + 128*math.sin(2*3.14*fr*i)

F1 = np.fft.fft2(im)

F2 = np.fft.fftshift(F1)

# Calculate a 2D magnitude spectrum
psd2D =np.abs(F2)
plt.figure(1)
plt.clf()
plt.imshow(im, cmap=plt.cm.gray)

plt.figure(2)
plt.clf()
plt.imshow(psd2D, extent=[-SZ/2,SZ/2, -SZ/2, SZ/2])

X,Y = np.meshgrid(X, Y)
fig = plt.figure(3)
plt.clf()
ax = Axes3D(fig)
ax.plot_surface(X, Y, im, rstride=1, cstride=1, cmap=cm.jet)
plt.show()

